#include "LabelEquivalenceSeg.h"
#include <stdlib.h>
#include <iostream>

#include <cuda_runtime.h>

LabelEquivalenceSeg::LabelEquivalenceSeg(int width, int height):
show(height, width){
	this->width = width;
	this->height = height;

	cudaMalloc(&img, sizeof(int) * width * height);
	cudaMalloc(&ref, sizeof(int) * width * height);
	cudaMalloc(&counter, sizeof(int) * width * height);
	cudaMalloc(&boundary_device, sizeof(bool) * width * height);


	color_pool = new unsigned char*[width * height];
	boundary_host = new bool[width * height];
	label = new int[width * height];
	
	for(int i = 0; i < width * height; i++){
		color_pool[i] = new unsigned char[3];
		color_pool[i][0] = (unsigned char)rand();
		color_pool[i][1] = (unsigned char)rand();
		color_pool[i][2] = (unsigned char)rand();
	}
}


LabelEquivalenceSeg::~LabelEquivalenceSeg(){
	cudaFree(img);
	cudaFree(ref);
	cudaFree(counter);
	show.~Mat_<cv::Vec3b>();
	delete[] color_pool;
	delete[] boundary_host;
	delete[] label;
}

void LabelEquivalenceSeg::viewSegmentResult(){
	int* temp = new int[width * height];
	cudaMemcpy(temp, img, width * height * sizeof(int), cudaMemcpyDeviceToHost);
	cudaMemcpy(boundary_host, boundary_device, sizeof(bool) * width * height, cudaMemcpyDeviceToHost);
	

	cv::Vec3b* pixel;
	cv::Vec3b* pixelSurf;
	unsigned char* color;
	for(int i = 0; i < height; i++){
		for(int j = 0; j < width; j++){
			pixel = &show.at<cv::Vec3b>(cv::Point2d(j, i));

			if(boundary_host[j + i * width] && temp[j + i * width] > -1){
				color = color_pool[temp[j + i * width]];
				pixel->val[0] = color[0];
				pixel->val[1] = color[1];
				pixel->val[2] = color[2];
			} else {
				pixel->val[0] = 0;
				pixel->val[1] = 0;
				pixel->val[2] = 0;
			}
		}
	}

	delete[] temp;	
	cv::imshow("Labled surface boundary", show);
	cv::waitKey(1);
}

int* LabelEquivalenceSeg::getLabelMap() const{
	return label;
}

bool* LabelEquivalenceSeg::getBoundaryMap() const{
	return boundary_host;
}
